%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                                                                         %
%                              Firm Moments                               %
%                                                                         %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


% ---------------------------- Description -------------------------------- 

    % This scripts picks the equilibrium path for the main variables of the
    % model to compute some firm moments along the transition. For an
    % overview of this part of the code, see Section (RP-3.3).

% ---------------------------------------------------------------------
%                     1. Firm Age Distrbution                         %
% ---------------------------------------------------------------------

    % Determines the number of periods after the shock takes place
    TimePeriods = length(p.TimeGrid);

    % Equilibrium path for the variety growth rate
    gN_vector = InfoMatrix(3,:);

    % Equilibrium path for the creative destruction rate
    tau_vector = p.alpha/(1-p.alpha)*(gN_vector + p.d0);

    % Equilibrium path for flow entry rate
    z_vector = 1/p.alpha*tau_vector-p.xstar;

    % Equilibrium path for the net product accumulation rate
    psi_vector =  p.xstar - (tau_vector+p.d0);

    % Preallocate a matrix to store the share of firms older than 10y
    Firms_10y_t = zeros(1, TimePeriods);

    % Preallocate a matrix to store markups by age along transition
    Markup_t = zeros(p.maxTime/p.periodlength + 1, TimePeriods);

    % Preallocate vector to store aggregate markups along transition
    AvgMup = zeros(1, TimePeriods);

    % We run a loop over possible age points to fill markup matrix
    for a = 1:(p.maxTime/p.periodlength + 1)

        % At every age point a, we interpolate between BGPs
        Markup_t(a, :) = linspace(CostMarkup_A(a), CostMarkup_A_counter(a), TimePeriods);

    end

    % For every time period...
    for t = 1:TimePeriods

        % Numerator of firm age density on period t - see Equation (RP-55)
        Firms_function_t = @(a) psi_vector(t)*exp((psi_vector(t)-gN_vector(t)).*a)./(psi_vector(t) - p.xstar*(1-exp(psi_vector(t).*a)));
    
        % Integrate over firm age for denominator - see Equation (RP-55)
        Firms_t = integral(Firms_function_t, 0, Inf);

        % Compute firm age density at time t  - see Equation (RP-55)
        Firms_function_t = @(a) Firms_function_t(a) / Firms_t;

        % Share of firms above 10 years -  - see Equation (RP-56)
        Firms_10y_t(t) = integral(Firms_function_t, 10, Inf);

        % Compute density-weighted markups by age - see Equation (RP-59)
        aux_func = @(a) Firms_function_t(a) .* interp1(avec, Markup_t(:, t), a, "linear", "extrap");

        % Sum over firm age for aggregate markups in t - see Equation (RP-59)
        AvgMup(t) = integral(aux_func, 0, Inf);

    end

% ---------------------------------------------------------------------
%                       2. Differential Equation                      %
% ---------------------------------------------------------------------

    % Create a matrix to store the V function over time - v_t(n)
    Vmatrix = zeros(p.maxprods, TimePeriods);

    % Initial condition for differential equation - see Equation (RP-54)
    Vmatrix(:,1) = ProductDist / (AvFirmSize/LFN);

    % For every time period...
    for t = 2:TimePeriods
        
        % For the first product, there cannot be a (n-1) term - see Equations (RP-53)
        dvdt = z_vector(t-1) + Vmatrix(2,t-1)*2*(tau_vector(t-1)+p.d0) - Vmatrix(1,t-1)*(p.xstar+tau_vector(t-1)+p.d0+gN_vector(t-1));
        Vmatrix(1,t) = Vmatrix(1,t-1) + dvdt*p.yearlength;

        % For the last product, there cannot be a (n+1) term - see Equations (RP-53)
        n = p.maxprods;
        dvdt = Vmatrix(n-1,t-1)*(n-1)*p.xstar-Vmatrix(n,t-1)*n*(p.xstar+tau_vector(t-1)+p.d0+gN_vector(t-1)/n);
        Vmatrix(n,t) = Vmatrix(n,t-1) + dvdt*p.yearlength;

        % For all other products, we compute the full differential equation - see Equations (RP-53)
        for n = 2:p.maxprods-1
        dvdt = Vmatrix(n-1,t-1)*(n-1)*p.xstar+Vmatrix(n+1,t-1)*(n+1)*(tau_vector(t-1)+p.d0)-Vmatrix(n,t-1)*n*(p.xstar+tau_vector(t-1)+p.d0+gN_vector(t-1)/n);
        Vmatrix(n,t) = Vmatrix(n,t-1) + dvdt*p.yearlength;
        end

    end

% ---------------------------------------------------------------------
%                          3. Firm Moments                            %
% ---------------------------------------------------------------------

    % Equilibrium path for the labor/variety ratio
    LFN_t = InfoMatrix(2,:);

    % Compute the path for the expected number of products - see Equation (RP-50)
    xProds_t = 1 ./ sum(Vmatrix, 1);

    % Average Firm Size equilibrium path - see Equation (RP-50)
    AvFirmSize_t = LFN_t .* xProds_t;

    % Entry Rate equilibrium path - see Equation (RP-51)
    EntryRate_t = z_vector .* xProds_t;

    % Exit Rate equilibrium path - see Equation (RP-52)
    ExitRate_t = (tau_vector + p.d0) .* Vmatrix(1,:) .* xProds_t;

    % Compute the share of firms with at least two products - see Equations (RP-57) and (RP-58)
    Prods2 = 1 - Vmatrix(1,:) .* xProds_t;

    % Compute the share of firms with at least five products - see Equations (RP-57) and (RP-58)
    Prods5 = sum(Vmatrix(5:end, :), 1) .* xProds_t;

 